import concurrent.futures
import queue
import json

import cv2
from flask import Flask, jsonify
from rapidocr import RapidOCR

from Invoice import extract_table_from_pdf
from checkImg import is_suitable_for_ocr
from readTravel import extract_trip_info
from readTrain import extract_traip_info

# app = Flask(__name__)

# 创建线程池，可根据服务器性能调整最大工作线程数
executor = concurrent.futures.ThreadPoolExecutor(max_workers=5)

# 初始化 RapidOCR
ocr = RapidOCR(
    params={
        "lang_det": "ch_mobile",  # ch_server
        "lang_rec": "ch_mobile",
        "text_score": 0.8,  # 文本置信度
        "min_height": 25,
        "with_onnx ": True,
    }
)  # text_score=0.8,det_unclip_ratio =1.8, use_gpu=False

request_queue = queue.Queue()


def get_ocr_result(image_path):
    # 执行 OCR
    result = ocr(image_path, use_det=True, use_cls=True, use_rec=True)

    # print(result.txts)
    # 整理识别结果
    recognized_text = "\n".join([line for line in result.txts])
    return recognized_text


def is_pdf(img: str):
    img = img.lower()
    # 检查文件路径的扩展名是否为.pdf
    return img.endswith(".pdf")


def get_file_name(pdf_file_path: str):
    return pdf_file_path.removesuffix(".pdf")


def get_type(pdf_file_name: str) -> str:
    parts = pdf_file_name.split("-")
    print(parts)
    print(len(parts) - 1)
    return parts[len(parts) - 1]


def get_pdf_file_path(pdf_file_path):
    type = get_type(get_file_name(pdf_file_path))
    if type is not None:
        if type == "发票":
            result_str = extract_table_from_pdf(pdf_file_path)
            print(result_str)
            return result_str
        elif type == "行程单":
            result_str = extract_trip_info(pdf_file_path)
            print(result_str)
            return result_str
        else:
            #火车票
            result_str = extract_traip_info(pdf_file_path)
            print(result_str)
            return result_str
    return None


def read(image_path: str):
    image = cv2.imread(image_path)
    if image_path is not None:
        if is_pdf(image_path):
            result = get_pdf_file_path(image_path)
        else:
            if is_suitable_for_ocr(image):
                # 将请求放入队列
                request_queue.put(image_path)
                img = request_queue.get()
                # 提交 OCR 任务到线程池
                future = executor.submit(get_ocr_result, img)
                result = future.result()
            else:
                return (
                    jsonify({"error": "图像可能不适合进行 OCR，请检查图像质量。"}),
                    200,
                )
        if result is not None:
            types = result["基本信息"]["类别"] if judge(str(result)) else None
            if types is None:
                type = result.__contains__("行程单")
                type2 = result.__contains__("发票")
                type3 = (
                    result.__contains__("站")
                    and result.__contains__("票")
                    and result.__contains__("报销")
                )
                if type:
                    types = "行程单"
                elif type2:
                    types = "发票"
                elif type3:
                    types = "火车票"
                else:
                    return jsonify({"error": "该类型不支持识别。"}), 200
            return jsonify({"text": result, "type": types}), 200
        else:
            return jsonify({"error": "图形可能不适合进行 OCR 识别,请检查图像质量。"}),200
    else:
        return jsonify({"error": "无法读取图像，请检查文件路径。"}), 200

def judge(result:str):
    result_str_fixed = result.replace("'", '"')
    result_dict = json.loads(result_str_fixed)
    basic_info = result_dict.get("基本信息", {})
    # 第二步：确保"基本信息"是字典，且包含"类别"键
    return isinstance(basic_info, dict) and "类别" in basic_info